import networkx as nx
import pandas as pd
from datetime import datetime, timedelta
import re
import json

class PathOptimizer:
    def __init__(self):
        self.graph = nx.DiGraph()
        self.city_mapping = {
            '北京': ['北京', '北京南', '北京西', '北京北', '北京东', '首都机场', '大兴机场'],
            '上海': ['上海', '上海虹桥', '上海南', '上海西', '浦东机场', '虹桥机场'],
            '广州': ['广州', '广州南', '广州东', '白云机场'],
            '深圳': ['深圳', '深圳北', '深圳东', '宝安机场'],
            '南京': ['南京', '南京南', '南京站', '禄口机场'],
            '杭州': ['杭州', '杭州东', '杭州南', '萧山机场'],
            '成都': ['成都', '成都东', '成都南', '双流机场', '天府机场'],
            '武汉': ['武汉', '武汉站', '汉口', '武昌', '天河机场'],
            '西安': ['西安', '西安北', '咸阳机场'],
            '重庆': ['重庆', '重庆北', '重庆西', '江北机场']
        }
    
    def parse_user_intent(self, user_message):
        """
        解析用户意图，提取关键信息
        """
        intent = {
            'type': 'unknown',
            'origin': None,
            'destination': None,
            'departure_date': None,
            'optimization_type': 'time'
        }
        
        # 简单的关键词匹配（后续可以集成更复杂的NLP）
        if any(keyword in user_message for keyword in ['去', '到', '从', '出发', '路线', '怎么走']):
            intent['type'] = 'route_search'
            
            # 提取城市信息
            for city, aliases in self.city_mapping.items():
                for alias in aliases:
                    if alias in user_message:
                        if '从' in user_message and user_message.find('从') < user_message.find(alias):
                            intent['origin'] = city
                        elif '到' in user_message and user_message.find('到') < user_message.find(alias):
                            intent['destination'] = city
                        elif not intent['origin']:
                            intent['origin'] = city
                        elif not intent['destination']:
                            intent['destination'] = city
            
            # 提取时间信息
            if '明天' in user_message:
                intent['departure_date'] = (datetime.now() + timedelta(days=1)).strftime('%Y-%m-%d')
            elif '后天' in user_message:
                intent['departure_date'] = (datetime.now() + timedelta(days=2)).strftime('%Y-%m-%d')
            elif '今天' in user_message:
                intent['departure_date'] = datetime.now().strftime('%Y-%m-%d')
            
            # 提取优化类型
            if any(keyword in user_message for keyword in ['便宜', '省钱', '价格']):
                intent['optimization_type'] = 'price'
            elif any(keyword in user_message for keyword in ['快', '时间', '早']):
                intent['optimization_type'] = 'time'
            elif any(keyword in user_message for keyword in ['舒适', '方便', '换乘']):
                intent['optimization_type'] = 'comfort'
        
        return intent
    
    def build_transport_graph(self, transport_data):
        """
        构建交通网络图
        """
        self.graph.clear()
        
        for route in transport_data:
            # 添加节点（站点+时间）
            origin_node = f"{route['origin']}_{route['departure_time']}"
            dest_node = f"{route['destination']}_{route['arrival_time']}"
            
            # 计算边的权重
            if route['type'] == 'high_speed_rail':
                time_weight = self._calculate_time_weight(route['departure_time'], route['arrival_time'])
                price_weight = route['price'] * 0.1  # 价格权重
                comfort_weight = 1  # 高铁舒适度较高
            elif route['type'] == 'flight':
                time_weight = self._calculate_time_weight(route['departure_time'], route['arrival_time'])
                price_weight = route['price'] * 0.1
                comfort_weight = 2  # 飞机舒适度中等
            else:  # 汽车等其他方式
                time_weight = self._calculate_time_weight(route['departure_time'], route['arrival_time'])
                price_weight = route['price'] * 0.1
                comfort_weight = 3  # 汽车舒适度较低
            
            # 添加边
            self.graph.add_edge(
                origin_node, dest_node,
                weight=time_weight + price_weight + comfort_weight,
                route_info=route
            )
    
    def _calculate_time_weight(self, departure_time, arrival_time):
        """
        计算时间权重
        """
        try:
            dep_time = datetime.strptime(departure_time, '%H:%M')
            arr_time = datetime.strptime(arrival_time, '%H:%M')
            
            # 处理跨天情况
            if arr_time < dep_time:
                arr_time += timedelta(days=1)
            
            duration = (arr_time - dep_time).total_seconds() / 3600  # 小时
            return duration * 2  # 时间权重系数
        except:
            return 24  # 默认权重
    
    def find_optimal_routes(self, origin, destination, departure_date, transport_data, optimization_type='time'):
        """
        寻找最优路径
        """
        # 构建交通图
        self.build_transport_graph(transport_data)
        
        # 寻找所有可能的路径
        all_routes = []
        
        # 获取所有可能的起点和终点
        origin_nodes = [node for node in self.graph.nodes() if origin in node]
        dest_nodes = [node for node in self.graph.nodes() if destination in node]
        
        for origin_node in origin_nodes:
            for dest_node in dest_nodes:
                try:
                    # 使用Dijkstra算法找最短路径
                    path = nx.shortest_path(self.graph, origin_node, dest_node, weight='weight')
                    
                    # 计算路径总信息
                    total_time = 0
                    total_price = 0
                    route_details = []
                    
                    for i in range(len(path) - 1):
                        edge_data = self.graph[path[i]][path[i+1]]
                        route_info = edge_data['route_info']
                        route_details.append(route_info)
                        
                        # 计算时间
                        time_diff = self._calculate_time_weight(
                            route_info['departure_time'], 
                            route_info['arrival_time']
                        )
                        total_time += time_diff
                        total_price += route_info['price']
                    
                    route_summary = {
                        'total_time': round(total_time, 1),
                        'total_price': total_price,
                        'transfers': len(route_details) - 1,
                        'details': route_details,
                        'score': self._calculate_route_score(total_time, total_price, len(route_details)-1, optimization_type)
                    }
                    
                    all_routes.append(route_summary)
                    
                except nx.NetworkXNoPath:
                    continue
        
        # 根据优化类型排序
        if optimization_type == 'time':
            all_routes.sort(key=lambda x: x['total_time'])
        elif optimization_type == 'price':
            all_routes.sort(key=lambda x: x['total_price'])
        elif optimization_type == 'comfort':
            all_routes.sort(key=lambda x: x['transfers'])
        else:
            all_routes.sort(key=lambda x: x['score'])
        
        return all_routes[:5]  # 返回前5条最优路线
    
    def _calculate_route_score(self, time, price, transfers, optimization_type):
        """
        计算路线综合评分
        """
        if optimization_type == 'time':
            return time * 0.6 + price * 0.2 + transfers * 0.2
        elif optimization_type == 'price':
            return price * 0.6 + time * 0.2 + transfers * 0.2
        else:  # comfort
            return transfers * 0.6 + time * 0.2 + price * 0.2
    
    def generate_route_response(self, routes, intent):
        """
        生成AI回复
        """
        if not routes:
            return "抱歉，没有找到合适的路线。请检查出发地和目的地是否正确。"
        
        response = f"为您找到从{intent['origin']}到{intent['destination']}的最优路线：\n\n"
        
        for i, route in enumerate(routes[:3], 1):
            response += f"方案{i}：\n"
            response += f"总时间：{route['total_time']}小时\n"
            response += f"总价格：{route['total_price']}元\n"
            response += f"换乘次数：{route['transfers']}次\n"
            
            response += "详细路线：\n"
            for j, detail in enumerate(route['details']):
                response += f"  {j+1}. {detail['type']} {detail['departure_time']}-{detail['arrival_time']} "
                response += f"{detail['origin']}→{detail['destination']} ({detail['price']}元)\n"
            
            response += "\n"
        
        return response